import sqlite3
import pandas as pd
from abc import ABC, abstractmethod
from typing import Any, List, Dict, Optional, Union, Tuple
import re
import numpy as np
from pydantic import BaseModel, Field, validator
from nl2sql.similarity.cosine_similarity import CosineSimilarity
from nl2sql.model.embedding_model import EmbeddingModel


class Enumeration(BaseModel):
    value: str  # 枚举值的值
    cor_tables: Dict[str, str]  # 对应的表名和列名，格式: {"table_name": "", "column_name": ""}
    vector: List[float] = []  # 对应着这个枚举值的向量

    def __str__(self) -> str:
        return f"枚举值为 {self.value}, 对应的表和列 {self.cor_tables}"


class Column(BaseModel):
    name: str = None  # 列名称
    description: str | None = None  # 列的描述（原desc）
    enums: List[Any] = []  # 列的枚举值


class Table(BaseModel):
    name: str = None  # 表名
    description: str | None = None  # 表的描述（原decs）
    columns: List[Column] = []  # 表的列


class BaseDatasource(ABC):
    """数据源基类，提供基本的数据操作功能"""

    def __init__(self,
                 embedding_model: Optional[EmbeddingModel] = None,
                 batch_size: int = 10):
        """
        初始化数据源

        Args:
            embedding_model (Optional[EmbeddingModel]): 嵌入模型，默认为None

            batch_size (int): 批处理大小，默认为10
        """
        self.name = None
        self.enumerations: List[Enumeration] = []
        self.data: Optional[pd.DataFrame] = None
        self.db_path: Optional[str] = None
        self.connection: Optional[sqlite3.Connection] = None
        self.embedding_model: Optional[EmbeddingModel] = embedding_model
        self.cosine_similarity: Optional[CosineSimilarity] = None

        self.batch_size: int = batch_size

        # 主要的存储
        self.tables: List[Table] = []  # 数据的信息

        if self.embedding_model:
            self.cosine_similarity = CosineSimilarity(embedding_model=embedding_model)

    @abstractmethod
    def load_data(self) -> None:
        """抽象方法，加载数据，子类必须实现"""
        pass

    def get_column_names(self) -> List[str]:
        """
        获取数据中所有的列名
        :return: 列名列表
        """
        return [column.name
                for table in self.tables
                for column in table.columns]

    @staticmethod
    def _filter_valid_enums(enums: List[Any]) -> List[str]:
        """筛选有效的枚举值，过滤掉无效值和数字类型"""
        if not enums:
            return []

        valid_enums = []
        for enum in enums:
            # 排除 None 值
            if enum is None:
                continue

            # 排除数字类型（整数、浮点数等）
            if isinstance(enum, (int, float)) and not isinstance(enum, bool):
                continue

            # 排除空字符串
            if isinstance(enum, str) and enum.strip() == '':
                continue

            # 将有效枚举值转换为字符串并添加到结果列表
            valid_enums.append(str(enum))

        return valid_enums

    def initialize_enums(self) -> None:
        """
        初始化枚举值的信息
        """
        if not self.embedding_model:
            print("警告: 嵌入模型未初始化，无法生成枚举值向量")
            return

        for table in self.tables:
            for column in table.columns:
                # 获取并筛选枚举值
                enums = self._filter_valid_enums(column.enums)
                if not enums:
                    continue

                # 分批处理枚举值
                for i in range(0, len(enums), self.batch_size):
                    batch_enums = enums[i:i + self.batch_size]
                    try:
                        batch_embeddings = self.embedding_model.get_text_embeddings(texts=batch_enums)
                        self._create_enumerations(batch_enums, batch_embeddings, table.name, column.name)
                    except Exception as e:
                        print(f"无法为枚举值生成向量 (批次 {i // self.batch_size + 1}): {e}")

    def _create_enumerations(self,
                             enums: List[str],
                             embeddings: List[List[float]],
                             table_name: str, column_name: str) -> None:
        """创建枚举值对象并添加到列表"""
        for enum, embedding in zip(enums, embeddings):
            cor_table = {
                "table_name": table_name,
                "column_name": column_name,
            }
            self.enumerations.append(Enumeration(
                value=enum,
                cor_tables=cor_table,
                vector=embedding
            ))

    def get_column_enums(self, column_name: str) -> List[Any]:
        """
        获取指定列的所有枚举值

        Args:
            column_name (str): 列名

        Returns:
            List[Any]: 该列的枚举值列表，如果列不存在则返回空列表

        Raises:
            ValueError: 如果列名在多个表中存在，提示需要指定表名
        """
        columns = [column for table in self.tables
                   for column in table.columns
                   if column.name == column_name]

        if not columns:
            print(f"⚠️ 列 '{column_name}' 不存在")
            return []
        if len(columns) > 1:
            raise ValueError(f"列名 '{column_name}' 在多个表中存在，请指定表名")

        return columns[0].enums

    def find_closest_enum(self,
                          input_value: str,
                          similarity_threshold: float) -> Optional[Dict[str, Any]]:
        """
        根据输入的不精确 enumerate value，找到最相近的枚举值及其对应的字段和数据表，只匹配有效文字类型的枚举。
        使用嵌入向量和 embedding_model 计算语义相似度。

        Args:
            input_value (str): 不精确的枚举值
            similarity_threshold

        Returns:
            Optional[Dict[str, Any]]: 匹配到的详细信息或 None
        """
        if not input_value or not isinstance(input_value, str):
            print(f"输入值无效: {input_value}")
            return None

        if not self.embedding_model or not self.enumerations:
            print("嵌入模型或枚举值未初始化")
            return None

        try:
            input_embedding = self.embedding_model.get_text_embeddings([input_value])[0]
        except Exception as e:
            print(f"为输入值生成嵌入向量失败: {input_value}, 错误: {e}")
            return None

        # 使用向量化操作计算所有相似度
        input_vec = np.array(input_embedding)
        enum_vectors = np.array([enum.vector for enum in self.enumerations if enum.vector])

        if not enum_vectors.size:
            print("没有可用的枚举向量")
            return None

        # 计算余弦相似度
        norms = np.linalg.norm(enum_vectors, axis=1) * np.linalg.norm(input_vec)
        norms = np.where(norms == 0, 1e-10, norms)  # 避免除零错误
        similarities = np.dot(enum_vectors, input_vec) / norms

        # 找到最佳匹配
        max_idx = np.argmax(similarities)
        highest_score = similarities[max_idx]

        if highest_score < similarity_threshold:
            print(f"⚠️ 最高相似度 {highest_score:.2f} 低于阈值 {similarity_threshold}")
            return None

        best_enum = next((enum for enum in self.enumerations if enum.vector and
                          np.array_equal(enum.vector, enum_vectors[max_idx])), None)

        if not best_enum:
            return None

        # 获取表和列的描述
        table_name = best_enum.cor_tables.get("table_name")
        column_name = best_enum.cor_tables.get("column_name")
        table_desc, column_desc = self._get_table_column_descriptions(table_name, column_name)

        return {
            "table_name": table_name,
            "table_description": table_desc,
            "column_name": column_name,
            "column_description": column_desc,
            "closest_enum": best_enum.value,
            "similarity_score": highest_score
        }

    def _get_table_column_descriptions(self, table_name: str, column_name: str) -> Tuple[Optional[str], Optional[str]]:
        """获取表和列的描述"""
        table_desc = None
        column_desc = None

        for table in self.tables:
            if table.name == table_name:
                table_desc = table.description
                for column in table.columns:
                    if column.name == column_name:
                        column_desc = column.description
                        break
                break

        return table_desc, column_desc

    def get_data_info(self) -> str:
        """
        获取数据的核心信息，专为大模型生成SQL设计

        Returns:
            Dict[str, Any]: 包含SQL生成所需的数据信息
        """
        info = {
            "tables": []
        }

        for table in self.tables:
            table_info = {
                "name": table.name,
                "description": table.description,
                "columns": []
            }

            for column in table.columns:
                column_info = {
                    "name": column.name,
                    "description": column.description,
                    "enum_values": self._filter_valid_enums(column.enums)
                }
                table_info["columns"].append(column_info)

            info["tables"].append(table_info)

        return str(info)

    def execute_sql(self, sql: str) -> pd.DataFrame:
        """
        执行SQL查询
        Args:
            sql (str): SQL查询语句
        Returns:
            pd.DataFrame: 查询结果
        """

        if not self.connection:
            raise ValueError("No database connection. Please create SQLite database first.")

        try:
            return pd.read_sql_query(sql, self.connection)
        except Exception as e:
            print(f"SQL执行错误: {e}")
            raise

    def close_connection(self) -> None:
        """关闭数据库连接"""
        if self.connection:
            try:
                self.connection.close()
                print("Database connection closed")
            except Exception as e:
                print(f"关闭数据库连接失败: {e}")
            finally:
                self.connection = None

    def __del__(self) -> None:
        """析构函数，确保关闭数据库连接"""
        self.close_connection()


if __name__ == "__main__":
    # 示例用法
    china_eastern = Enumeration(value='中国东方航空', cor_tables={"table_name": "airlines", "column_name": "carrier"})
    print(china_eastern)
